import struct
import pickle
import socket
import time

class IOClosedException(BaseException):
  pass

def int_to_bytes(intval):
  return struct.pack("<i", intval)

def bytes_to_int(byteval):
  return struct.unpack("<i", byteval)[0]

def get_len(sock, data, length):
  if (data != None and len(data) >= length):
    return data[0:length], data[length:]
  else:
    while True:
      try:
        next_data = sock.recv(4096)
        break
      except socket.error as e:
        if e.errno == 11:
          time.sleep(0.01)
          continue

    if(len(next_data) == 0):
      raise IOClosedException()
    if(data != None):
      data += next_data
    else:
      data = next_data
    return get_len(sock, data, length)

def get_string(sock, data = None):
  length_str, data = get_len(sock, data, 4)
  return get_len(sock, data, bytes_to_int(length_str))

def to_len_prefix_string(bytes_string):
  return int_to_bytes(len(bytes_string)) + bytes_string

class Message(object):
  def __init__(self, name, args):
    self.name = name
    self.args = args

  def write_to(self, sock):
    data = (to_len_prefix_string(bytes(self.name, "utf8")) +
        to_len_prefix_string(pickle.dumps(self.args)))
    sock.send(data)

class Reply(object):
  def __init__(self, name, args):
    self.name = name
    self.args = args

  def __call__(self, *args):
    return Message(self.name, args)

class Command(object):
  def __init__(self, name, args):
    self.name = name
    self.args = args
    self.replies = dict()

  def reply(self, name, *args):
    self.replies[name] = Reply(name, args)
    return self

  def get_reply(self, name):
    if name in self.replies:
      return self.replies[name]

  def __call__(self, sock, *args):
    msg = Message(self.name, args)
    msg.write_to(sock)
    name, data = get_string(sock)
    pickle_string, data = get_string(sock, data)
    args = pickle.loads(pickle_string)
    return Message(name, args)

class Schema(object):
  def __init__(self):
    self.cmds = dict()

  def __getattr__(self, key):
    if key in self.cmds:
      return self.cmds[key]
    return getattr(super(), key)

  def cmd(self, name, *args):
    cmd = Command(name, args)
    self.cmds[name] = cmd
    return cmd

  def handle(self, dispatcher):
    try:
      name, data = get_string(dispatcher)
      name = str(name, "utf8")
      pickle_string, data = get_string(dispatcher, data)
      args = pickle.loads(pickle_string)
      context = Context(self.cmds[name])
      reply = getattr(dispatcher, name)(context, *args)
      reply.write_to(dispatcher)
    except IOClosedException:
      dispatcher.close()



class Context(object):
  def __init__(self, cmd):
    self.cmd = cmd

  def __getattr__(self, key):
    reply = self.cmd.get_reply(key)
    if reply:
      return reply
    return getattr(super(), key)

